from decimal import Decimal

import pandas as pd
import numpy as np

from app.services.energy_dispatch.energy_dispatch_simple.energy_dispatch_simple import calculate_wind_output


# --------------------------
# 1. 环境定义（强化供电不足惩罚）
# --------------------------
class MicroGridEnv:
    def __init__(self, env_df, config):
        self.df = env_df
        self.config = config
        self.reset()

    def reset(self):
        self.current_step = 0
        self.battery = self.config["初始电量(kWh)"]
        self.total_cost = 0.0
        self.total_carbon = 0.0
        self.power_shortage_count = 0  # 供电不足计数器
        self.step_actions = []
        self.charge_discharge_count = 0  # 充放电次数
        self.prev_battery_power = 0  # 上一步电池功率，用于检测充放电切换
        return self._get_state()

    def _get_state(self):
        row = self.df.iloc[self.current_step]
        max_vals = [1000, 20, 50, 2, 1, 1, self.config["电池容量(kWh)"]]

        return np.array([
            row["太阳辐照度(W/m²)"] / max_vals[0],
            row["风速(m/s)"] / max_vals[1],
            row["负荷功率(kW)"] / max_vals[2],
            row["电价_国网(元/kWh)"] / max_vals[3],
            row["电价_风能(元/kWh)"] / max_vals[4],
            row["电价_太阳能(元/kWh)"] / max_vals[5],
            self.battery / max_vals[6]
        ])

    def step(self, action):
        charge_flag, power, grid_purchase = action
        row = self.df.iloc[self.current_step]
        load = row["负荷功率(kW)"]

        # 计算可再生能源出力
        # solar_output = max(0, row["太阳辐照度(W/m²)"] * 0.01)
        # wind_output = max(0, row["风速(m/s)"] * 0.5)
        # total_renewable = solar_output + wind_output

        # 参数
        solar_area = float("500.0")              # 光伏面积 m²
        solar_efficiency = float("0.20")         # 光伏效率
        wind_turbine_count = 3                     # 风机数量
        single_wind_power = float("30.0")        # 单台风机额定功率 kW

        # 获取环境数据
        solar_irradiance = float(row["太阳辐照度(W/m²)"])
        wind_speed = float(str(row["风速(m/s)"]))   # 注意转换为 Decimal 类型

        # 光伏出力
        solar_output = (solar_irradiance * solar_area * solar_efficiency) / float("1000")  # kW

        # 风电出力（调用你的函数 × 风机数量）
        wind_output = float(calculate_wind_output(Decimal(wind_speed), Decimal(single_wind_power)) * wind_turbine_count)

        # 总可再生能源出力
        total_renewable = solar_output + wind_output
        # 处理电池动作（含效率）
        input_power = 0  # 充电输入（含损耗）
        output_power = 0  # 放电输出（含损耗）
        battery_power = 0  # 电池功率（正值表示充电，负值表示放电）
        if charge_flag == 1:  # 充电
            max_possible_charge = min(
                self.config["最大充放电功率(kW)"],
                self.config["电池容量(kWh)"] - self.battery
            )
            input_power = min(power, max_possible_charge / self.config["充电效率"])
            charge_power = input_power * self.config["充电效率"]
            self.battery += charge_power
            battery_power = input_power  # 正值表示充电
        else:  # 放电
            max_possible_discharge = min(
                self.config["最大充放电功率(kW)"],
                self.battery
            )
            discharge_power = min(power, max_possible_discharge)
            output_power = discharge_power * self.config["放电效率"]
            self.battery -= discharge_power
            battery_power = -discharge_power  # 负值表示放电

        # 检测充放电状态切换
        if (battery_power > 0 and self.prev_battery_power <= 0) or (battery_power < 0 and self.prev_battery_power >= 0):
            self.charge_discharge_count += 1
        self.prev_battery_power = battery_power

        # 能量平衡计算（强化供电保障）
        shortage = 0
        if charge_flag == 1:
            total_energy_needed = load + input_power  # 负荷+充电需求
            total_energy_available = total_renewable + grid_purchase
            if total_energy_available < total_energy_needed:
                shortage = total_energy_needed - total_energy_available
                grid_purchase += shortage  # 强制补充购电
                self.power_shortage_count += 1
        else:
            total_supply = total_renewable + grid_purchase + output_power  # 负荷供给
            if total_supply < load:
                shortage = load - total_supply
                grid_purchase += shortage  # 强制补充购电
                self.power_shortage_count += 1

        # 计算成本与碳排放
        cost = grid_purchase * row["电价_国网(元/kWh)"]
        carbon = grid_purchase * self.config["碳排放系数(kg/kWh)"]
        self.total_cost += cost
        self.total_carbon += carbon

        # 记录动作
        self.step_actions.append({
            "时间步": self.current_step,
            "电池操作": "充电" if charge_flag == 1 else "放电",
            "电池功率(kW)": battery_power,
            "电池电量(kWh)": self.battery,
            "购电量(kW)": grid_purchase,
            "太阳能出力(kW)": solar_output,
            "风能出力(kW)": wind_output,
            "负荷功率(kW)": load,
            "供电缺口(kW)": shortage,
            "成本(元)": cost,
            "碳排放(kg)": carbon
        })

        # 优化奖励函数：大幅提高供电不足惩罚，降低成本/碳排权重
        shortage_penalty = self.power_shortage_count * 200  # 惩罚从20提高到200
        reward = - (cost * 0.5 + carbon * 0.05) - shortage_penalty  # 降低成本/碳排权重，优先保供电
        if shortage > 0:
            reward -= shortage * 50  # 对每次缺口额外惩罚

        # 推进时间步
        self.current_step += 1
        done = self.current_step >= len(self.df)
        next_state = self._get_state() if not done else None

        return next_state, reward, done, {"成本": cost, "碳排放": carbon, "缺口": shortage}

    def get_action_history(self):
        rows = []
        for step in self.step_actions:
            solar = step["太阳能出力(kW)"]
            wind = step["风能出力(kW)"]
            grid = step["购电量(kW)"]
            battery = step["电池电量(kWh)"]
            battery_capacity = self.config["电池容量(kWh)"]

            # 价格信息来自 env_df 中每一时刻
            row_df = self.df.iloc[step["时间步"]]
            price_grid = row_df["电价_国网(元/kWh)"]
            price_wind = row_df["电价_风能(元/kWh)"]
            price_solar = row_df["电价_太阳能(元/kWh)"]

            # 单小时步长
            dt = 1.0

            # 成本计算
            solar_cost = solar * price_solar * dt
            wind_cost = wind * price_wind * dt
            grid_cost = grid * price_grid * dt

            # 碳排计算
            solar_emission = solar * 0.02 * dt
            wind_emission = wind * 0.05 * dt
            grid_emission = grid * self.config["碳排放系数(kg/kWh)"] * dt

            # 总供电
            total_power = solar + wind + grid

            # 电池使用率
            battery_power = abs(step["电池功率(kW)"])
            battery_usage = battery_power * dt
            battery_usage_rate = (battery_usage / battery_capacity * 100) if battery_capacity > 0 else 0.0

            rows.append({
                "time_point": step["时间步"],
                "battery_energy": battery,
                "battery_usage_rate": battery_usage_rate,
                "solar_power": solar,
                "wind_power": wind,
                "grid_power": grid,
                "total_power": total_power,
                "load": step["负荷功率(kW)"],
                "solar_cost": solar_cost,
                "wind_cost": wind_cost,
                "grid_cost": grid_cost,
                "solar_emission": solar_emission,
                "wind_emission": wind_emission,
                "grid_emission": grid_emission,
            })

        return pd.DataFrame(rows)

    def calculate_battery_metrics(self):
        """计算电池使用的详细指标"""
        if not self.step_actions:
            return {
                "平均电池电量百分比": 0,
                "充放电次数": 0,
                "最大电池电量(kWh)": 0,
                "最小电池电量(kWh)": 0,
                "深度放电次数": 0
            }

        history_df = pd.DataFrame(self.step_actions)
        if history_df.empty:
            return {
                "平均电池电量百分比": 0,
                "充放电次数": 0,
                "最大电池电量(kWh)": 0,
                "最小电池电量(kWh)": 0,
                "深度放电次数": 0
            }

        # 计算平均电池电量百分比
        avg_battery_level = history_df["电池电量(kWh)"].mean()
        avg_battery_percentage = (avg_battery_level / self.config["电池容量(kWh)"]) * 100

        # 计算充放电次数（已在step中统计）
        charge_discharge_count = self.charge_discharge_count

        # 计算最大/最小电池电量
        max_battery = history_df["电池电量(kWh)"].max()
        min_battery = history_df["电池电量(kWh)"].min()

        # 计算深度放电次数（电量低于20%）
        depth_discharge_threshold = self.config["电池容量(kWh)"] * 0.2
        depth_discharge_count = (history_df["电池电量(kWh)"] < depth_discharge_threshold).sum()

        return {
            "平均电池电量百分比": avg_battery_percentage,
            "充放电次数": charge_discharge_count,
            "最大电池电量(kWh)": max_battery,
            "最小电池电量(kWh)": min_battery,
            "深度放电次数": depth_discharge_count
        }
